// -*- C++ -*-
#ifndef V3D_GPU_KLT_H
#define V3D_GPU_KLT_H

# if defined(V3DLIB_GPGPU_ENABLE_CG)

#include "v3d_gpubase.h"
#include "v3d_gpupyramid.h"

namespace V3D_GPU
{

   struct KLT_TrackerBase
   {
         KLT_TrackerBase(int nIterations, int nLevels, int levelSkip, int windowWidth)
            : _nIterations(nIterations), _nLevels(nLevels),
              _levelSkip((levelSkip > 0) ? levelSkip : (nLevels-1)),
              _windowWidth(windowWidth),
              _width(0), _height(0), _featureWidth(0), _featureHeight(0),
              _margin(5), _convergenceThreshold(0.01), _SSD_Threshold(1000.0)
         { }

         ~KLT_TrackerBase() { }

         void setBorderMargin(float margin) { _margin = margin; }
         void setConvergenceThreshold(float thr) { _convergenceThreshold = thr; }
         void setSSD_Threshold(float thr) { _SSD_Threshold = thr; }

         void swapFeatureBuffers() { std::swap(_featuresBuffer0, _featuresBuffer1); }

         unsigned int sourceFeatureTexID() const { return _featuresBuffer0->textureID(); }
         unsigned int targetFeatureTexID() const { return _featuresBuffer1->textureID(); }

         void activateTargetFeatures() { _featuresBuffer1->activate(); }

      protected:
         int const _nIterations, _nLevels, _levelSkip, _windowWidth;
         int       _width, _height, _featureWidth, _featureHeight;
         float     _margin, _convergenceThreshold, _SSD_Threshold;

         RTT_Buffer *_featuresBuffer0, *_featuresBuffer1;
   }; // end struct KLT_TrackerBase

//----------------------------------------------------------------------

   struct KLT_Tracker : public KLT_TrackerBase
   {
         KLT_Tracker(int nIterations = 5, int nLevels = 4, int levelSkip = -1, int windowWidth = 7)
            : KLT_TrackerBase(nIterations, nLevels, levelSkip, windowWidth),
              _featuresBufferA("rgb=32f", "KLT_Tracker::_featuresBufferA"),
              _featuresBufferB("rgb=32f", "KLT_Tracker::_featuresBufferB"),
              _trackingShader(0)
         {
            _featuresBuffer0 = &_featuresBufferA;
            _featuresBuffer1 = &_featuresBufferB;
         }

         ~KLT_Tracker() { }

         void allocate(int width, int height, int featureWidth, int featureHeight);
         void deallocate();

         void provideFeatures(float const * features);
         void readFeatures(float * features);
         void trackFeatures(unsigned int pyrTex0, unsigned int pyrTex1);

      protected:
         RTT_Buffer _featuresBufferA, _featuresBufferB;

         Cg_FragmentProgram * _trackingShader;
   }; // end struct KLT_Tracker

//----------------------------------------------------------------------

   struct KLT_TrackerWithGain : public KLT_TrackerBase
   {
         KLT_TrackerWithGain(int nIterations = 5, int nLevels = 4, int levelSkip = -1, int windowWidth = 7)
            : KLT_TrackerBase(nIterations, nLevels, levelSkip, windowWidth),
              _featuresBufferA("rgb=32f", "KLT_TrackerWithGain::_featuresBufferA"),
              _featuresBufferB("rgb=32f", "KLT_TrackerWithGain::_featuresBufferB"),
              _featuresBufferC("rgb=32f", "KLT_TrackerWithGain::_featuresBufferC"),
              _trackingShader(0)
         {
            _featuresBuffer0 = &_featuresBufferA;
            _featuresBuffer1 = &_featuresBufferB;
            _featuresBuffer2 = &_featuresBufferC;
         }

         ~KLT_TrackerWithGain() { }

         void allocate(int width, int height, int featureWidth, int featureHeight);
         void deallocate();

         void provideFeaturesAndGain(float const * features);
         void readFeaturesAndGain(float * features);
         void trackFeaturesAndGain(unsigned int pyrTex0, unsigned int pyrTex1);

      protected:
         RTT_Buffer _featuresBufferA, _featuresBufferB, _featuresBufferC;

         RTT_Buffer * _featuresBuffer2;

         Cg_FragmentProgram * _trackingShader;
   }; // end struct KLT_TrackerWithGain

//----------------------------------------------------------------------

   struct KLT_Detector
   {
         KLT_Detector(int minDist = 5)
            : _minDist(minDist), _margin(10.0f),
              _convRowsBuffer("rgb=32f", "KLT_Detector::_convRowsBuffer"),
              _cornernessBuffer("rgba=8", "KLT_Detector::_cornernessBuffer"),
              _nonmaxRowsBuffer("rgba=8", "KLT_Detector::_nonmaxRowsBuffer"),
              _pointListBuffer("rgb=32f", "KLT_Detector::_pointListBuffer"),
              _cornernessShader1(0), _cornernessShader2(0), _nonmaxShader(0), _traverseShader(0)
         { }

         ~KLT_Detector() { }

         void setBorderMargin(float margin) { _margin = margin; }

         void allocate(int width, int height, int pointListWidth, int pointListHeight);
         void deallocate();

         // Note: the texID should be a texture as generated by GPUPyramidWithDerivativesCreator.
         void detectCorners(float const minCornerness, unsigned int texID, int& nFeatures,
                            int nPresentFeatures = 0, float * presentFeaturesBuffer = 0);
         void extractCorners(int const nFeatures, float * dest);

      protected:
         int const _minDist;

         int _width, _height;
         int _pointListWidth, _pointListHeight;
         int _histpyrWidth, _nHistpyrLevels;

         float _margin;

         RTT_Buffer _convRowsBuffer; // temporary buffer for structure matrix accumulation
         RTT_Buffer _cornernessBuffer;
         RTT_Buffer _nonmaxRowsBuffer; // temporary buffer for non-max suppression
         RTT_Buffer _pointListBuffer;

         unsigned int _histpyrTexID;
         unsigned int _histpyrFbIDs[16];
         unsigned int _vbo;

         Cg_FragmentProgram * _cornernessShader1;
         Cg_FragmentProgram * _cornernessShader2;
         Cg_FragmentProgram * _nonmaxShader;
         Cg_FragmentProgram * _traverseShader;
   }; // end struct KLT_Detector

   struct KLT_TrackedFeature
   {
         KLT_TrackedFeature()
            : status(-1), gain(1.0f)
         { }

         //! 0 means tracked from previous frame, 1 is newly created and -1 means invalidated track.
         int status;
         float pos[2];
         float gain;
   };

//----------------------------------------------------------------------

   struct KLT_SequenceTrackerConfig
   {
         KLT_SequenceTrackerConfig()
            : nIterations(5), //!< #Iterations in the tracker
              nLevels(4), //!< Levels used in the coarse-to-fine approach
              levelSkip(3), //!< Levels skipped, reasonable values are 1 (all levels) and nLevels-1 (2 levels)
              windowWidth(7), //!< Size of the floating window
              trackBorderMargin(10.0f), //!< In pixels
              convergenceThreshold(0.02f), //!< In pixels
              SSD_Threshold(1000.0f),
              trackWithGain(false),
              minDistance(10), //!< Radius of non-max feature suppression
              minCornerness(10.0f),
              detectBorderMargin(10.0f) //!< In pixels
         { }

         int   nIterations, nLevels, levelSkip, windowWidth;
         float trackBorderMargin, convergenceThreshold, SSD_Threshold;
         bool  trackWithGain;

         int   minDistance;
         float minCornerness, detectBorderMargin;
   }; // end struct KLT_SequenceTrackerConfig

   struct KLT_SequenceTracker
   {
         KLT_SequenceTracker(KLT_SequenceTrackerConfig const& config)
            : _config(config), _trackWithGain(config.trackWithGain),
              _detector(config.minDistance),
              _tracker(0), _trackerWithGain(0)
         {
            _pyrCreator0 = &_pyrCreatorA;
            _pyrCreator1 = &_pyrCreatorB;
         }

         ~KLT_SequenceTracker()
         { }

         void allocate(int width, int height, int nLevels,
                       int featuresWidth, int featuresHeight)
         {
            this->allocate(width, height, nLevels, featuresWidth, featuresHeight, 2*featuresWidth, 2*featuresHeight);
         }
         void allocate(int width, int height, int nLevels,
                       int featuresWidth, int featuresHeight,
                       int pointListWidth, int pointListHeight);
         void deallocate();

         void setBorderMargin(float margin)
         {
            if (_tracker)
               _tracker->setBorderMargin(margin);
            else
               _trackerWithGain->setBorderMargin(margin);

            _detector.setBorderMargin(margin);
         }

         void setConvergenceThreshold(float thr)
         {
            if (_tracker)
               _tracker->setConvergenceThreshold(thr);
            else
               _trackerWithGain->setConvergenceThreshold(thr);
         }

         void setSSD_Threshold(float thr)
         {
            if (_tracker)
               _tracker->setSSD_Threshold(thr);
            else
               _trackerWithGain->setSSD_Threshold(thr);
         }

         void detect(unsigned char const * image, int& nDetectedFeatures, KLT_TrackedFeature * dest);
         void redetect(unsigned char const * image, int& nNewFeatures, KLT_TrackedFeature * dest);
         void track(unsigned char const * image, int& nPresentFeatures, KLT_TrackedFeature * dest);

         void advanceFrame()
         {
            if (_tracker)
               _tracker->swapFeatureBuffers();
            else
               _trackerWithGain->swapFeatureBuffers();

            std::swap(_pyrCreator0, _pyrCreator1);
         }

         unsigned int getCurrentFrameTextureID() const { return _pyrCreator1->sourceTextureID(); }

      protected:
         struct FeaturePoint
         {
               float data[3];

               bool operator<(FeaturePoint const& b) const
               {
                  return this->data[2] > b.data[2];
               }
         }; // end struct FeaturePoint

         KLT_SequenceTrackerConfig const _config;

         bool const _trackWithGain;

         int _width, _height;
         int _featuresWidth, _featuresHeight;
         int _pointListWidth, _pointListHeight;

         PyramidWithDerivativesCreator _pyrCreatorA;
         PyramidWithDerivativesCreator _pyrCreatorB;
         KLT_Detector _detector;

         KLT_Tracker         * _tracker;
         KLT_TrackerWithGain * _trackerWithGain;

         PyramidWithDerivativesCreator * _pyrCreator0;
         PyramidWithDerivativesCreator * _pyrCreator1;

         FeaturePoint * _corners;
   }; // end struct KLT_SequenceTracker

} // end namespace V3D_GPU

# endif // defined(V3DLIB_GPGPU_ENABLE_CG)

#endif
